import os
import sys
from matplotlib import pyplot as plt
import numpy as np

WORK_DIR = os.path.dirname(os.path.abspath(__file__))
os.chdir(WORK_DIR)

sys.path.insert(0, WORK_DIR)
# WORK_DIR = os.path.join(WORK_DIR, "elbow")


from video_composite import images_load

def data_load(trajs_path, visib_path, queries_path=None, H=None, W=None):
    
    if queries_path is not None:
        query_points = np.load(queries_path)
        # W, H = query_points[0].shape
        query_points = (query_points*np.array([W,H])).astype(np.uint32)
        print(f"query point sample: W X H = {query_points[0]}")
        # print(f"query_points: {query_points.shape}")
        # print(f"query_point: {query_points[724]}")
    trajs = np.load(trajs_path)[0]
    visibies = np.load(visib_path)[0]
    print(f"len of trajs: {len(trajs)}, number of points: {len(trajs[0])}, ndim of points: {trajs.shape[-1]}")
    print(f"visibility of points: {visibies.shape}")
    if queries_path is not None:
        return query_points, trajs, visibies
    
    return trajs, visibies

def extract_one_point_traj(obj_dir="elbow", coord=724, trajs=None, visibies=None, 
                           cam_rid=None, cam_traj=None, cam_visibies=None, 
                           save_traj=False, perfix=None):
    point_traj = []
    visib_chain = []
    unvisib_times = []
    visib_frag = 1
    unvisib_start = -1
    unvisib_ids = []
    
    for frame, traj, visib in zip(range(len(trajs)), trajs, visibies):
        # print(f"frame: {frame}, traj: {traj[coord]}, visib: {visib[coord]}")
        point_traj.append(traj[coord])
        visib_chain.append(visib[coord][0])
        if not visib[coord][0]:
            unvisib_ids.append(frame)
            
        if visib_frag and not visib[coord][0]:
            unvisib_start = frame - 1
            visib_frag = 0
        elif not visib_frag and visib[coord][0]:
            unvisib_times.append((unvisib_start, frame))
            unvisib_start = -1
            visib_frag = 1
    point_traj, unvisib_times = np.array(point_traj), np.array(unvisib_times)
    print(f"unvisib: {len(unvisib_ids)}: {unvisib_ids}")
    
    if save_traj:
        _elbow_dir = os.path.join(WORK_DIR, obj_dir)
        save_path = os.path.join(_elbow_dir, f"{perfix}_track_{str(coord)}_trajectory.txt")
        np.savetxt(save_path, point_traj, fmt='%d')
        save_path = os.path.join(_elbow_dir, f"{perfix}_track_{str(coord)}_visibility.txt")
        np.savetxt(save_path, visib_chain, fmt='%d')
    
    if cam_traj is not None:
        cam_track_traj = []
        cam_track_unvisib = []
        for frame, trj, visb in zip(range(len(cam_traj)), cam_traj, cam_visibies):
            cam_track_traj.append(trj[cam_rid])
            if not visb[cam_rid][0]:
                cam_track_unvisib.append(frame)

        print(f"cam_track_unvisib: {len(cam_track_unvisib)}: {cam_track_unvisib}")
        if cam_track_traj:
            cam_track_traj = np.array(cam_track_traj)
            _elbow_dir = os.path.join(WORK_DIR, obj_dir)
            save_path = os.path.join(_elbow_dir, f"{perfix}_cama_track_{str(cam_rid)}_trajectory.txt")
            np.savetxt(save_path, cam_track_traj, fmt='%d')
            save_path = os.path.join(_elbow_dir, f"{perfix}_cama_track_{str(cam_rid)}_visibility.txt")
            np.savetxt(save_path, cam_track_unvisib, fmt='%d')
    
    return (point_traj, visib_chain, unvisib_times)

def camara_motion_compensate(point_traj, cam_track_traj, save_file=None):
    # point_traj = point_traj[:, [1,0]]
    # cam_track_traj = cam_track_traj[:, [1,0]]
    
    # 1. 计算相机运动补偿
    camara_offset_traj = cam_track_traj - cam_track_traj[0]
    # 2. 计算相机运动补偿后的轨迹
    point_traj_compensate = point_traj - camara_offset_traj
    # 3. 保存补偿后的轨迹
    # 规定左上角为坐标原点绘图
    
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    save_path = os.path.join(_elbow_dir, f"{save_file}_camara_compenated.txt")
    np.savetxt(save_path, point_traj_compensate, fmt='%d')
    # save_path = os.path.join(_elbow_dir, f"{save_file}_camara_offset.txt")
    # np.savetxt(save_path, camara_offset_traj, fmt='%d')
    
    # plt.scatter(camara_offset_traj[:,0], camara_offset_traj[:,1], marker='.', label="camara")
    # plt.plot(point_traj[:,0], point_traj[:,1], marker='.', label="point")       
    # plt.plot(point_traj_compensate[:,0], point_traj_compensate[:,1], marker='.', label="point_compensate")
    # plt.scatter(camara_offset_traj[0,0], camara_offset_traj[0,1], marker='.', label="cam_start")
    # plt.scatter(camara_offset_traj[-1,0], camara_offset_traj[-1,1], marker='.', label="cam_end")
    # plt.scatter(point_traj_compensate[0,0], point_traj_compensate[0,1], marker='o', label="point_start")
    # plt.scatter(point_traj_compensate[-1,0], point_traj_compensate[-1,1], marker='*', lw=3, label="point_end")
    # # 调整坐标轴方向，使得左上角为原点
    # plt.gca().invert_yaxis()
    # plt.legend()
    # plt.show()
    # pass


def test_camara_motion_compensate():
    traj_path = "results_add_goe_feats_track_724_trajectory.txt" # W,H
    cam_track_traj_path = "results_add_goe_feats_cama_track_293_trajectory.txt" # W,H
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    traj = np.loadtxt(os.path.join(_elbow_dir, traj_path))
    cam_track_traj = np.loadtxt(os.path.join(_elbow_dir, cam_track_traj_path))
    
    camara_motion_compensate(traj, cam_track_traj, save_file=traj_path.split(".")[0])


def right_hand_compensate(point_traj, right_hand_traj, save_file=None):
    # 1. 计算右手运动补偿
    right_hand_offset_traj = right_hand_traj - right_hand_traj[0]
    # 2. 计算右手运动补偿后的轨迹
    point_traj_compensate = point_traj - right_hand_offset_traj
    # 3. 保存补偿后的轨迹
    plt.scatter(right_hand_offset_traj[:,0], right_hand_offset_traj[:,1], marker='.', label="right_hand_offset")
    plt.plot(point_traj[:,0], point_traj[:,1], marker='.', label="point")       
    plt.plot(point_traj_compensate[:,0], point_traj_compensate[:,1], marker='.', label="point_compensate")
    plt.scatter(right_hand_offset_traj[0,0], right_hand_offset_traj[0,1], marker='.', label="right_start")
    plt.scatter(right_hand_offset_traj[-1,0], right_hand_offset_traj[-1,1], marker='.', label="right_end")
    plt.scatter(point_traj_compensate[0,0], point_traj_compensate[0,1], marker='o', label="point_start")
    plt.scatter(point_traj_compensate[-1,0], point_traj_compensate[-1,1], marker='*', lw=3, label="point_end")
    # 调整坐标轴方向，使得左上角为原点
    plt.gca().invert_yaxis()
    plt.legend()
    plt.show()
    
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    save_path = os.path.join(_elbow_dir, f"{save_file}_right_hand_compenated.txt")
    np.savetxt(save_path, point_traj_compensate, fmt='%d')

def test_right_hand_compensate():
    traj_path = "results_add_goe_feats_track_724_trajectory_camara_compenated.txt" # W,H
    
    # right_hand_traj, right_hand_visib_chain, _ = extract_one_point_traj_wrapper(
    #     coord_id=373, traj_path_dir="results", obj_dir="right_thumb", save_traj=True )
    right_thumb_path = os.path.join(WORK_DIR, "right_thumb", "results_track_373_trajectory.txt")
    right_hand_traj = np.loadtxt(right_thumb_path)
    
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    traj = np.loadtxt(os.path.join(_elbow_dir, traj_path))

    right_hand_compensate(traj, right_hand_traj, save_file=traj_path.split(".")[0])


def single_traj_compensate(traj_path="track_724_trajectory.txt", 
                           visib_path="track_724_visibility.txt", 
                           save_traj=False):
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    traj = np.loadtxt(os.path.join(_elbow_dir, traj_path))
    visib = np.loadtxt(os.path.join(_elbow_dir, visib_path))
    
    xMax, xMin = traj[:, 0].max(), traj[:, 0].min()
    yMax, yMin = traj[:, 1].max(), traj[:, 1].min()
    
    unvisib_times = []
    visib_frag = 1
    unvisib_start = -1
    for frame,v in enumerate(visib):        
        if visib_frag and not v:
            unvisib_start = frame
            visib_frag = 0
        elif not visib_frag and v:
            unvisib_times.append((unvisib_start, frame - 1))
            visib_frag = 1
    unvisib_times = np.array(unvisib_times)    
    
    for start, end in unvisib_times:
        num_interpolation = end-start+1
        # linear interpolation
        # _compoint = np.linspace(traj[start], traj[end], num_interpolation)
        
        # 抛物线插值        
        x, y = np.vstack((traj[start-2], traj[start-1], traj[end], traj[end+1])).T
        coefficients = np.polyfit(x, y, 2)
        
        nx = np.linspace(traj[start-1, 0], traj[end+1, 0], num_interpolation)
        ny = np.array([np.polyval(coefficients, x) for x in nx])
        ny = np.clip(ny, yMin, yMax)
        _compoint = np.vstack((nx, ny)).T
        
        # replace the unvisib part with the interpolated part
        traj[start:end+1] = _compoint
    
    # # 维度对调
    # traj = traj[:,[1,0]] # X:H, Y:W
    
    if save_traj:
        _elbow_dir = os.path.join(WORK_DIR, "elbow")
        save_path = os.path.join(_elbow_dir, f"{traj_path.split('.')[0]}_compensate.txt")
        np.savetxt(save_path, traj, fmt='%d')
    
    return traj

def check_traj_continuity(traj, jump_scale=2.5):
    # check the continuity of the trajectory
    tj_diff = np.diff(traj, axis=0)
    tj_diff = np.insert(tj_diff, 0, 0, axis=0)
    
    diff_norm = np.linalg.norm(tj_diff, axis=1)
    diff_max = np.max(diff_norm)
    diff_min = np.min(diff_norm)
    diff_mean = np.mean(diff_norm)
    diff_std = np.std(diff_norm)
    point_jumps = np.where(diff_norm>diff_mean*jump_scale)[0]
    
    print(f"diff_max: {diff_max}, diff_min: {diff_min}, diff_mean: {diff_mean}, diff_std: {diff_std}")
    print(f"more than mean coords: {point_jumps}, diff_norm: {diff_norm[point_jumps]}")
    
    return {
        'diff_max': diff_max,
        'diff_min': diff_min,
        'diff_mean': diff_mean,
        'diff_std': diff_std,
        'point_jumps': point_jumps
    }

def test_draw_traj(traj):
    from draw_trajectory import draw_trajectory_2d, draw_trajectory_2d_animation

    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    if isinstance(traj, (str or os.path.normpath)):
        print(f"load traj from {os.path.join(_elbow_dir, traj)}")
        traj = np.loadtxt(os.path.join(_elbow_dir, traj))
    
    
    result = check_traj_continuity(traj)
    point_jumps = result['point_jumps']
    
    draw_trajectory_2d(traj, jumps=point_jumps, save_gif=1, filename=os.path.join(_elbow_dir, "trajectory_2D.jpg"))
    draw_trajectory_2d_animation(traj, jumps=point_jumps, save_gif=0)

def extract_one_point_traj_wrapper(coord_id, traj_path_dir, obj_dir="elbow", save_traj=False):
    _elbow_dir = os.path.join(WORK_DIR, obj_dir)
    traj_path_dir = os.path.join(_elbow_dir, traj_path_dir)
    
    traj_path = os.path.join(traj_path_dir, "trajectory.npy")
    visib_path = os.path.join(traj_path_dir, "visibility.npy")
    
    trajs, visibies = data_load(trajs_path=traj_path, visib_path=visib_path)
    (point_traj, visib_chain, unvisib_times) = extract_one_point_traj(obj_dir=obj_dir, coord=coord_id, trajs=trajs, visibies=visibies, 
        save_traj=save_traj, perfix=traj_path_dir)
    
    return (point_traj, visib_chain, unvisib_times)

def test_extract_one_point_traj():
    _elbow_dir = os.path.join(WORK_DIR, "elbow")
    queries_path=os.path.join(_elbow_dir, "sam2_mask/query_dense.npy")
    trajs_path=os.path.join(_elbow_dir, "results_only_optical_flow/trajectory.npy")
    visib_path=os.path.join(_elbow_dir, "results_only_optical_flow/visibility.npy")
    
    trajs_path=os.path.join(_elbow_dir, "results_add_keypoints/trajectory.npy")
    visib_path=os.path.join(_elbow_dir, "results_add_keypoints/visibility.npy")
    
    trajs_path=os.path.join(_elbow_dir, "results_add_goe_feats/trajectory.npy")
    visib_path=os.path.join(_elbow_dir, "results_add_goe_feats/visibility.npy")
    
    cam_trajs_path=os.path.join("./camara", "results/trajectory.npy")
    cam_visib_path=os.path.join("./camara", "results/visibility.npy")
    cam_queries_path=os.path.join("./camara", "sam2_mask/query_dense.npy")
    
    
    
    image_data = images_load(images_path=_elbow_dir)
    H,W = image_data.images.shape[-3:-1]
    
    perfix = trajs_path.split('/')[-2]
    
    query_points, trajs, visibies = data_load(
        queries_path=queries_path, trajs_path=trajs_path, visib_path=visib_path, H=H, W=W)
    cam_query_points, cam_trajs, cam_visibies = data_load(
        queries_path=cam_queries_path, trajs_path=cam_trajs_path, visib_path=cam_visib_path, H=H, W=W)


    track_coord_id=724
    camara_ref_id=293
    point_traj, visib_chain, unvisib_times = extract_one_point_traj(
        coord=track_coord_id, trajs=trajs, visibies=visibies, 
        cam_rid=camara_ref_id, cam_traj=cam_trajs, cam_visibies=cam_visibies, 
        save_traj=1, perfix=perfix)


if __name__ == '__main__':
    print()
    # traj_path = "results_only_optical_flow_track_724_trajectory.txt"
    # visib_path = "results_only_optical_flow_track_724_visibility.txt"

    test_right_hand_compensate()
    # test_extract_one_point_traj()
    # test_camara_motion_compensate()
    # single_traj_compensate(traj_path, visib_path, save_traj=1)
    # test_draw_traj(traj="track_724_trajectory_compensate.txt")
    
